package jia.tools;

import jia.game.GenericFlags;
import org.apache.log4j.Logger;
import org.apache.xml.security.encryption.EncryptedData;
import org.apache.xml.security.encryption.EncryptedKey;
import org.apache.xml.security.encryption.XMLCipher;
import org.apache.xml.security.keys.KeyInfo;
import org.w3c.dom.Document;
import org.w3c.dom.Element;

import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.xml.transform.OutputKeys;
import javax.xml.transform.Transformer;
import javax.xml.transform.TransformerFactory;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.security.Key;

public class EncryptTool{
    private static final Logger log = Logger.getLogger(EncryptTool.class);
    static{
        org.apache.xml.security.Init.init();
    }
    
    public static Document parseFile(String fileName) throws Exception{
        javax.xml.parsers.DocumentBuilderFactory dbf = javax.xml.parsers.DocumentBuilderFactory.newInstance();
        dbf.setNamespaceAware(true);
        javax.xml.parsers.DocumentBuilder db = dbf.newDocumentBuilder();
        Document document = db.parse(fileName);
        return document;
    }
    
    private static SecretKey GenerateKeyEncryptionKey() throws Exception{
        String jceAlgorithmName = "DESede";
        KeyGenerator keyGenerator = KeyGenerator.getInstance(jceAlgorithmName);
        SecretKey keyEncryptKey = keyGenerator.generateKey();
        return keyEncryptKey;
    }
    
    private static void storeKeyFile(Key keyEncryptKey) throws IOException{
        byte[] keyBytes = keyEncryptKey.getEncoded();
        File keyEncryptKeyFile = new File("resources/config/keyEncryptKey");
        FileOutputStream outStream = new FileOutputStream(keyEncryptKeyFile);
        outStream.write(keyBytes);
        outStream.close();
        if (GenericFlags.debugMessages){
            log.debug("Key encryption key stored in: " + keyEncryptKeyFile.toURL().toString());
        }
    }
    
    private static SecretKey GenerateSymmetricKey() throws Exception{
        String jceAlgorithmName = "AES";
        KeyGenerator keyGenerator = KeyGenerator.getInstance(jceAlgorithmName);
        keyGenerator.init(128);
        return keyGenerator.generateKey();
    }
    
    private static void writeEncryptedDocToFile(Document doc, String fileName) throws Exception{
        File encryptionFile = new File(fileName);
        FileOutputStream outStream = new FileOutputStream(encryptionFile);
        TransformerFactory factory = TransformerFactory.newInstance();
        Transformer transformer = factory.newTransformer();
        transformer.setOutputProperty(OutputKeys.OMIT_XML_DECLARATION, "no");
        DOMSource source = new DOMSource(doc);
        StreamResult result = new StreamResult(outStream);
        transformer.transform(source, result);
        outStream.close();
        if (GenericFlags.debugMessages){
            log.debug("Encrypted XML document written to: " + encryptionFile.toURL().toString());
        }
    }
    
    /**
     * TODO:  remove this method
     **/
    private static void usage(){
        System.err.println("usage - java EncryptTool " + "infilename outfilename elementtoencrypt");
        System.err.println("example - java EncryptTool " + "test.xml encrypted.xml CreditCardNumber");
    }
    
    /**
     * TODO:  remove this method
     **/
    public static void main(String args[]) throws Exception{
        String[] testArgs = new String[2];
        testArgs[0] = "resources/levels/plaintext.xml";
        testArgs[1] = "resources/levels/output.xml";
        TestRun(testArgs);
    }
    
    
    /**
     * TODO:  remove this method
     **/
    private static void TestRun(String[] args) throws Exception{
        if(args.length < 2) {
            usage();
            System.exit(1);
        }
        
        // parse file into document
        Document document = parseFile(args[0]);
        //Document document = new ();
        
        // generate symmetric key
        Key symmetricKey = GenerateSymmetricKey();
        
        // Get a key to be used for encrypting the symmetric key
        Key keyEncryptKey = GenerateKeyEncryptionKey();
        
        // Write the key to a file
        storeKeyFile(keyEncryptKey);
        
        // initialize cipher
        XMLCipher keyCipher = XMLCipher.getInstance(XMLCipher.TRIPLEDES_KeyWrap);
        keyCipher.init(XMLCipher.WRAP_MODE, keyEncryptKey);
        
        // encrypt symmetric key
        EncryptedKey encryptedKey = keyCipher.encryptKey(document, symmetricKey);
        
        // specify the element to encrypt
        Element rootElement = document.getDocumentElement();
        Element elementToEncrypt = rootElement;
        if(args.length > 2) {
            elementToEncrypt = (Element) rootElement.getElementsByTagName(args[2]).item(0);
            if(elementToEncrypt == null) {
                System.err.println("Unable to find element: " + args[2]);
                System.exit(1);
            }
        }
        
        // initialize cipher
        XMLCipher xmlCipher = XMLCipher.getInstance(XMLCipher.AES_128);
        xmlCipher.init(XMLCipher.ENCRYPT_MODE, symmetricKey);
        
        // add key info to encrypted data element
        EncryptedData encryptedDataElement = xmlCipher.getEncryptedData();
        KeyInfo keyInfo = new KeyInfo(document);
        keyInfo.add(encryptedKey);
        encryptedDataElement.setKeyInfo(keyInfo);
        
        // do the actual encryption
        boolean encryptContentsOnly = false;
        xmlCipher.doFinal(document, elementToEncrypt, encryptContentsOnly);
        
        // write the results to a file
        writeEncryptedDocToFile(document, args[1]);
    }
    
    
    public static void designWriter(String xmlBlock, String fileName){
        
    }
    
    
    /**
     * Public method that will take a formatted level document (in XML) and encrypt it, while
     * finally writing it to a seperate file.
     * @param document
     * @param fileName
     * @param destFileName
     */
    
    public static void levelWriter(Document document, String fileName, String destFileName){
        try {
            
            String destFilePath = "resources/levels/" + destFileName + ".xml";
            // generate symmetric key
            Key symmetricKey = GenerateSymmetricKey();
            
            // Get a key to be used for encrypting the symmetric key
            Key keyEncryptKey = GenerateKeyEncryptionKey();
            
            // Write the key to a file
            storeKeyFile(keyEncryptKey);
            
            // initialize cipher
            XMLCipher keyCipher = XMLCipher.getInstance(XMLCipher.TRIPLEDES_KeyWrap);
            keyCipher.init(XMLCipher.WRAP_MODE, keyEncryptKey);
            
            // encrypt symmetric key
            EncryptedKey encryptedKey = keyCipher.encryptKey(document, symmetricKey);
            
            // specify the element to encrypt
            Element rootElement = document.getDocumentElement();
            Element elementToEncrypt = rootElement;
            
            // initialize cipher
            XMLCipher xmlCipher = XMLCipher.getInstance(XMLCipher.AES_128);
            xmlCipher.init(XMLCipher.ENCRYPT_MODE, symmetricKey);
            
            // add key info to encrypted data element
            EncryptedData encryptedDataElement = xmlCipher.getEncryptedData();
            KeyInfo keyInfo = new KeyInfo(document);
            keyInfo.add(encryptedKey);
            encryptedDataElement.setKeyInfo(keyInfo);
            
            // do the actual encryption
            boolean encryptContentsOnly = false;
            xmlCipher.doFinal(document, elementToEncrypt, encryptContentsOnly);
            
            // write the results to a file
            writeEncryptedDocToFile(document, destFilePath);
        } catch(Exception e) {
            if (GenericFlags.errorMessages){
                log.error(e.getMessage());
            }
        }
    }
}
